{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eacbb70c",
   "metadata": {},
   "source": [
    "# Lesson 5: Coding and Financial Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "babbc472",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5613ba9-d387-4b7d-9670-ec2b6e4353be",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "llm_config = {\"model\": \"gpt-4-turbo\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4e8fd85",
   "metadata": {},
   "source": [
    "## Define a code executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb90a672-4318-47de-a061-384e68dd75c9",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "from autogen.coding import LocalCommandLineCodeExecutor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5cafee7-a0ac-4583-8201-5e19eb603138",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "executor = LocalCommandLineCodeExecutor(\n",
    "    timeout=60,\n",
    "    work_dir=\"coding\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b62a1a0",
   "metadata": {},
   "source": [
    "## Create agents "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d9b3d6-9915-4fea-a853-512dfb77c4a5",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent, AssistantAgent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b741b947",
   "metadata": {},
   "source": [
    "### 1. Agent with code executor configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d286e5-b7b7-405a-b1d7-212c0cc84497",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "code_executor_agent = ConversableAgent(\n",
    "    name=\"code_executor_agent\",\n",
    "    llm_config=False,\n",
    "    code_execution_config={\"executor\": executor},\n",
    "    human_input_mode=\"ALWAYS\",\n",
    "    default_auto_reply=\n",
    "    \"Please continue. If everything is done, reply 'TERMINATE'.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1eb1094",
   "metadata": {},
   "source": [
    "### 2. Agent with code writing capability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c38c33-d335-48cd-96eb-cafe9ac70142",
   "metadata": {
    "height": 114
   },
   "outputs": [],
   "source": [
    "code_writer_agent = AssistantAgent(\n",
    "    name=\"code_writer_agent\",\n",
    "    llm_config=llm_config,\n",
    "    code_execution_config=False,\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fb1a62f-8dd2-4636-9103-b26ea27a29bd",
   "metadata": {
    "height": 46
   },
   "outputs": [],
   "source": [
    "code_writer_agent_system_message = code_writer_agent.system_message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c8e9a04-55fb-4124-a9d0-3431274476d2",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "print(code_writer_agent_system_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f487b75a",
   "metadata": {},
   "source": [
    "## The task!\n",
    "\n",
    "Ask the two agents to collaborate on a stock analysis task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cce0ba4-9fea-4d36-b4a0-a1e51fcd09a6",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "today = datetime.datetime.now().date()\n",
    "message = f\"Today is {today}. \"\\\n",
    "\"Create a plot showing stock gain YTD for NVDA and TLSA. \"\\\n",
    "\"Make sure the code is in markdown code block and save the figure\"\\\n",
    "\" to a file ytd_stock_gains.png.\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d02df6c",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#ECECEC; padding:15px; \"> <b>Note:</b> In this lesson, you will use GPT 4 for better results. Please note that the lesson has a quota limit. If you want to explore the code in this lesson further, we recommend trying it locally with your own API key."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77fc5f1b",
   "metadata": {},
   "source": [
    "**Note**: You might see a different set of outputs than those shown in the video. The agents collaborate to generate the code needed for your task, and they might produce code with errors in the process. However, they will ultimately provide a correct code in the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230bb9c0-9c80-498f-a05f-96cef11ddb4e",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "chat_result = code_executor_agent.initiate_chat(\n",
    "    code_writer_agent,\n",
    "    message=message,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b0cbbdd",
   "metadata": {},
   "source": [
    "## Let's see the plot!\n",
    "\n",
    "**Note**: \n",
    "* Your plot might differ from the one shown in the video because the LLM's freestyle code generation could choose a different plot type, such as a bar plot. \n",
    "* You can re-run the previous cell and check the generated code. If it produces a bar plot, remember you can directly specify your preference by asking for a specific plot type instead of a bar plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26804313-566f-42fb-9243-4e0c147f770b",
   "metadata": {
    "height": 80
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from IPython.display import Image\n",
    "\n",
    "Image(os.path.join(\"coding\", \"ytd_stock_gains.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccf49d7e",
   "metadata": {},
   "source": [
    "**Note**: The agent will automatically save the code in a .py file and the plot in a .png file. To access and check the files generated by the agents, go to the `File` menu and select `Open....` Then, open the folder named `coding` to find all the generated files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aded8dc",
   "metadata": {},
   "source": [
    "## User-Defined Functions\n",
    "\n",
    "Instead of asking LLM to generate the code for downloading stock data \n",
    "and plotting charts each time, you can define functions for these two tasks and have LLM call these functions in the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bca089a7-d85a-40ad-b33d-6a5318076c6e",
   "metadata": {
    "height": 403
   },
   "outputs": [],
   "source": [
    "def get_stock_prices(stock_symbols, start_date, end_date):\n",
    "    \"\"\"Get the stock prices for the given stock symbols between\n",
    "    the start and end dates.\n",
    "\n",
    "    Args:\n",
    "        stock_symbols (str or list): The stock symbols to get the\n",
    "        prices for.\n",
    "        start_date (str): The start date in the format \n",
    "        'YYYY-MM-DD'.\n",
    "        end_date (str): The end date in the format 'YYYY-MM-DD'.\n",
    "    \n",
    "    Returns:\n",
    "        pandas.DataFrame: The stock prices for the given stock\n",
    "        symbols indexed by date, with one column per stock \n",
    "        symbol.\n",
    "    \"\"\"\n",
    "    import yfinance\n",
    "\n",
    "    stock_data = yfinance.download(\n",
    "        stock_symbols, start=start_date, end=end_date\n",
    "    )\n",
    "    return stock_data.get(\"Close\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dff50a6-39ae-46a4-9b9c-e98c1550472b",
   "metadata": {
    "height": 352
   },
   "outputs": [],
   "source": [
    "def plot_stock_prices(stock_prices, filename):\n",
    "    \"\"\"Plot the stock prices for the given stock symbols.\n",
    "\n",
    "    Args:\n",
    "        stock_prices (pandas.DataFrame): The stock prices for the \n",
    "        given stock symbols.\n",
    "    \"\"\"\n",
    "    import matplotlib.pyplot as plt\n",
    "\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    for column in stock_prices.columns:\n",
    "        plt.plot(\n",
    "            stock_prices.index, stock_prices[column], label=column\n",
    "                )\n",
    "    plt.title(\"Stock Prices\")\n",
    "    plt.xlabel(\"Date\")\n",
    "    plt.ylabel(\"Price\")\n",
    "    plt.grid(True)\n",
    "    plt.savefig(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ff7175",
   "metadata": {},
   "source": [
    "### Create a new executor with the user-defined functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "452d315f-681d-4418-9cd0-653cc8d6a668",
   "metadata": {
    "height": 97
   },
   "outputs": [],
   "source": [
    "executor = LocalCommandLineCodeExecutor(\n",
    "    timeout=60,\n",
    "    work_dir=\"coding\",\n",
    "    functions=[get_stock_prices, plot_stock_prices],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b95b30ad-1e6d-4d88-9a6c-2b0f11b625e4",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "code_writer_agent_system_message += executor.format_functions_for_prompt()\n",
    "print(code_writer_agent_system_message)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d3cc77",
   "metadata": {},
   "source": [
    "### Let's update the agents with the new system message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "615c02b0-0175-442e-bed2-52bb35b8945c",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "code_writer_agent = ConversableAgent(\n",
    "    name=\"code_writer_agent\",\n",
    "    system_message=code_writer_agent_system_message,\n",
    "    llm_config=llm_config,\n",
    "    code_execution_config=False,\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9a725b0-9bb6-47e3-85df-27d9aadc9290",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "code_executor_agent = ConversableAgent(\n",
    "    name=\"code_executor_agent\",\n",
    "    llm_config=False,\n",
    "    code_execution_config={\"executor\": executor},\n",
    "    human_input_mode=\"ALWAYS\",\n",
    "    default_auto_reply=\n",
    "    \"Please continue. If everything is done, reply 'TERMINATE'.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8d7b57c",
   "metadata": {},
   "source": [
    "### Start the same task again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8895bb51-bcdd-4f63-a49f-858fc25c474f",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "chat_result = code_executor_agent.initiate_chat(\n",
    "    code_writer_agent,\n",
    "    message=f\"Today is {today}.\"\n",
    "    \"Download the stock prices YTD for NVDA and TSLA and create\"\n",
    "    \"a plot. Make sure the code is in markdown code block and \"\n",
    "    \"save the figure to a file stock_prices_YTD_plot.png.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c4bb718",
   "metadata": {},
   "source": [
    "### Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a930c007-dd56-4da6-b542-902312ed27dd",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "Image(os.path.join(\"coding\", \"stock_prices_YTD_plot.png\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c65bac5e",
   "metadata": {},
   "source": [
    "**Note**: The agent will automatically save the code in a .py file and the plot in a .png file. To access and check the files generated by the agents, go to the `File` menu and select `Open....` Then, open the folder named `coding` to find all the generated files."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
